[V1] V1 engine implements parallel sampling (AsyncLLM and LLMEngine) #10980

afeldman-nm · 2024-12-07T17:49:54Z

This PR adds support for parallel sampling to v1 AsyncLLM and LLMEngine.

Parallel sampling is implemented outside the engine. This does not impact the vLLM v0 parallel sampling implementation.

Design doc: https://docs.google.com/document/d/1_fvbHVCfexkPAj2Vx0Q0gNvE-b53WwtrLrD6NldgBFU/edit?usp=sharing (message me if you need access permissions)

A request with n>1 will spawn n requests with n=1 and aggregate their outputs in accordance with the output_kind. ~~If prefix caching is enabled, an initial warmup request with max_tokens=1 will be sent to the engine to fill the prefix cache.~~ Update: The vLLM v1 engine can exploit APC when a prompt repeats within a batch, even if that prompt was not seen in a previous batch. Therefore, no warmup request is required.

The abstractions are cleanest for async v1 engine because AsyncLLM presents a generate() method that handles adding requests and running the engine; parallel sampling is implemented by writing a wrapper which invokes this method n times against parallel sampling child requests.

v1 LLMEngine presents add_request() and step() methods, so parallel sampling is implemented by writing an add_request() wrapper which branches parallel sampling requests into n add_request() calls for child requests, and then having step() aggregate child request outputs into a parent request output.

FIX #13419
FIX #13420

github-actions · 2024-12-07T17:50:05Z

👋 Hi! Thank you for contributing to the vLLM project.
Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which starts running only a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of those by going to your fastcheck build on Buildkite UI (linked in the PR checks section) and unblock them. If you do not have permission to unblock, ping simon-mo or khluu to add you in our Buildkite org.

Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can do one of these:

Add ready label to the PR
Enable auto-merge.

🚀

mergify · 2024-12-13T15:47:04Z

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @afeldman-nm.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

FrederickVu · 2025-01-30T23:57:27Z

It seems like this PR is implementing ideas similar to those implemented in PR #9302 for the V0 engine. That PR created some issues that were addressed in PR #11898 and which may exist in the proposed V1 code.

In particular, the proposed code currently does not properly handle the case when a seed value is provided for the parent request; the seed value is duplicated in child requests, leading to identical outputs in the child requests. The fix in #11898 was simply to move the copying of the SamplingParams object inside a for-loop and to increment the seed value of the parent request.

Additionally, the proposed code for the V1 engine defines get_num_unfinished_requests() in a way that is currently incompatible with the same function for the V0 engine (though it's really the V0 engine's code which is in the wrong, I am choosing to mention this here as it seems relevant and the PR is still open). In the V0 engine, that function actually counts the number of SequenceGroup objects in the engine as opposed to the number of requests. This led to bugs mentioned in Issues #10949 and #11519, and an attempt to fix them is in PR #12428. However, the proposed fix is not compatible with the proposed V1 code. I will comment on that PR as well so that maintainers can adapt the V0 code to align with the proposed V1 code.

mergify · 2025-02-05T06:44:49Z

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @afeldman-nm.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

Signed-off-by: Andrew Feldman <afeldman@neuralmagic.com>

m-harmonic · 2025-02-14T17:56:42Z

@m-harmonic

Thanks for working on this. I haven't had a chance to look into the specifics of the new implementation but also wanted to ask about n>1 behavior as this PR (#9302) reduced our throughput by about 3x. It seems that sequence forking was extremely important in the v0 engine. Is an equivalent strategy going to be used in v1? We're now waiting for n>1 support in the v1 engine before we can upgrade past 0.6.3. Thanks

Thanks for the feedback. I think this is because you didn't add --enable-prefix-caching. Starting with the PR you mentioned, vLLM relies on prefix caching to deduplicate the prompts in parallel sampling. As prefix caching is not enabled by default in V0, you may see slowdown if you didn't add --enable-prefix-caching. With V1, prefix caching is enabled by default, so you should be able to enjoy the speedup without any extra config.

Hi thanks for the reply. We did pass enable_prefix_caching=True in the engine arguments, which did not affect the throughput. Is there something else we should be doing?
Edit: I got a notification for a comment by @ afeldman-nm but oddly enough it's not showing up now. Per your question, we see the decreased throughput in VLLM's log stats. The number of generated tokens / sec falls from roughly 7k/sec to 2k/sec. In contrast, the prompt throughput rises by about 4-5x. This is for a value of n between 50 and 100. We did have prefix caching enabled in the engine args.

Hi @m-harmonic , sorry I deleted my response because I thought @WoosukKwon 's suggestion to use prefix caching was more relevant, although it appears that that did not resolve the issue. Although, based on the metrics you shared, it sounds like prefix caching is working but the decode throughput is the issue.

Thanks @afeldman-nm, appreciate the reply. One thing I'm wondering is whether matching n>1 throughput performance pre-0.6.4 is a priority for this change or follow-up work. It seems like fundamentally this new approach confers lower throughput. Relatedly, is there an architectural reason why this new approach is necessary in v1 or is it possible forking could be added later on if it makes a big difference?

Signed-off-by: Andrew Feldman <afeldman@neuralmagic.com>

afeldman-nm · 2025-02-18T21:38:15Z

Merged #13421 (LLMEngine support for parallel sampling) into this PR.

Signed-off-by: Andrew Feldman <afeldman@neuralmagic.com>

njhill

Thanks @afeldman-nm for the great work.

I am still reviewing the llm_engine.py changes but wanted to post the comments I have so far. I feel like there may still be room to simplify/condense things a bit more, and possibly move more logic out of the main classes.

njhill · 2025-02-18T22:50:23Z

vllm/v1/engine/async_llm.py

+        async for out in merge_async_iterators(*gens):
+            yield out[1]  # out[0] is index


Suggested change

async for out in merge_async_iterators(*gens):

yield out[1] # out[0] is index

async for _, out in merge_async_iterators(*gens):

yield out

njhill · 2025-02-18T22:56:18Z

vllm/v1/engine/async_llm.py

@@ -241,6 +244,56 @@ async def generate(
            await self.abort(request_id)
            raise

+    async def _generate_parallel_sampling(


I'm wondering whether it might be cleaner to also move this method into ParallelSamplingRequestManager, but it takes the base generate` method as an additional arg.

njhill · 2025-02-18T22:57:34Z

vllm/v1/engine/async_llm.py

+    ) -> AsyncGenerator[RequestOutput, None]:
+        """Generate completions for parallel sampling requests."""
+        req_mgr = ParallelSamplingRequestManager(request_id, sampling_params)
+        n = req_mgr.n


nti: no need for this variable

njhill · 2025-02-18T22:58:06Z

vllm/v1/engine/async_llm.py

+        # Aggregate generators for n child requests
+        gens: List[AsyncGenerator[RequestOutput, None]] = []
+        for idx in range(n):
+            c_sampling_params = req_mgr.get_child_sampling_params(idx)


How about a single method that returns both the child sampling params and request id?

njhill · 2025-02-18T23:00:45Z

vllm/v1/engine/parallel_sampling.py

+          Child `sampling_params` instance.
+        """
+        seed = self.sampling_params.seed
+        if seed is None and self.cached_child_sampling_params:


nit:

Suggested change

if seed is None and self.cached_child_sampling_params:

if self.cached_child_sampling_params:

njhill · 2025-02-18T23:09:07Z

vllm/v1/engine/parallel_sampling.py

+            # Note: will be sorted by index later
+            self.request_output.outputs.append(new_completion)
+
+    def _get_parent_request_output(self) -> RequestOutput:


suggested rename:

Suggested change

def _get_parent_request_output(self) -> RequestOutput:

def _get_final_request_output(self) -> RequestOutput:

njhill · 2025-02-18T23:14:16Z

vllm/v1/engine/parallel_sampling.py

+from vllm.sampling_params import RequestOutputKind, SamplingParams
+
+
+class ParallelSamplingRequestManager:


WDYT about renaming this to

Suggested change

class ParallelSamplingRequestManager:

class ParallelSamplingRequest:

njhill · 2025-02-18T23:30:18Z

vllm/v1/engine/llm_engine.py

+    def _num_parallel_sampling_requests(self) -> int:
+        return len(self.parallel_parent_reqs)
+
+    def _num_parallel_sampling_child_requests(self) -> int:
+        return len(self.parallel_child_reqs)


IMO these helper methods are unnecessary (just adds more LOC)

njhill · 2025-02-18T23:30:49Z

vllm/v1/engine/llm_engine.py

+        self._do_reset_parallel_sampling = (True if num_parallel_reqs > 0 else
+                                            self._do_reset_parallel_sampling)


Suggested change

self._do_reset_parallel_sampling = (True if num_parallel_reqs > 0 else

self._do_reset_parallel_sampling)

self._do_reset_parallel_sampling |= num_parallel_reqs > 0

Actually, looking again, this would be clearer:

Suggested change

self._do_reset_parallel_sampling = (True if num_parallel_reqs > 0 else

self._do_reset_parallel_sampling)

if self.parallel_parent_requests:

self._do_reset_parallel_sampling = True

njhill · 2025-02-18T23:31:13Z

vllm/v1/engine/llm_engine.py

+        if num_parallel_reqs > 0 and len(request_outputs) > 0:
+            # Process parallel sampling child request outputs
+            return self._aggregate_parallel_sampling_outputs(request_outputs)
+        else:


nit: redundant else

njhill · 2025-02-19T00:46:51Z

vllm/v1/engine/llm_engine.py

+        self._do_reset_parallel_sampling = (True if num_parallel_reqs > 0 else
+                                            self._do_reset_parallel_sampling)


Actually, looking again, this would be clearer:

Suggested change

self._do_reset_parallel_sampling = (True if num_parallel_reqs > 0 else

self._do_reset_parallel_sampling)

if self.parallel_parent_requests:

self._do_reset_parallel_sampling = True

njhill · 2025-02-19T01:24:46Z

vllm/v1/engine/llm_engine.py

+                            prompt_adapter_request=prompt_adapter_request,
+                            priority=priority)
+
+    def _add_request_parallel_sampling(


Similar thought to the async case, perhaps this could be moved into the other class? (along with the state ... could have a subclass with that just for the sync case)

mergify bot added the frontend label Dec 7, 2024

mergify bot added the needs-rebase label Dec 13, 2024

afeldman-nm force-pushed the afeldman-nm/sample branch from 7cd9a24 to f506458 Compare December 16, 2024 16:05

mergify bot removed the needs-rebase label Dec 16, 2024

robertgshaw2-redhat mentioned this pull request Jan 30, 2025

[V1] Feedback Thread #12568

Open

FrederickVu mentioned this pull request Jan 31, 2025

[Bugfix] Fix tqdm progress bar when SamplingParams.n > 1 #12428

Open

mergify bot added the v1 label Feb 5, 2025

mergify bot added the needs-rebase label Feb 5, 2025

afeldman-nm added 3 commits February 9, 2025 04:59

wip

57f3329

Signed-off-by: Andrew Feldman <afeldman@neuralmagic.com>

wip

50584f6

Signed-off-by: Andrew Feldman <afeldman@neuralmagic.com>

Merge branch 'main' into sample

bf3cfd0

afeldman-nm force-pushed the afeldman-nm/sample branch from 3d5b962 to bf3cfd0 Compare February 9, 2025 05:47

mergify bot removed the needs-rebase label Feb 9, 2025

afeldman-nm added 3 commits February 10, 2025 15:18

stream=false works

98726ed

Signed-off-by: Andrew Feldman <afeldman@neuralmagic.com>

streaming nearly works

cd649df

Signed-off-by: Andrew Feldman <afeldman@neuralmagic.com>

Merge branch 'sample' into sample_merge

fdc3296

mergify bot added documentation Improvements or additions to documentation ci/build labels Feb 10, 2025

afeldman-nm added 4 commits February 11, 2025 13:10

refactor

a5415ef

Signed-off-by: Andrew Feldman <afeldman@neuralmagic.com>

Merge branch 'main' into sample_merge

eb9042a

good async implementation

a6637a9

Signed-off-by: Andrew Feldman <afeldman@neuralmagic.com>

Merge branch 'main' into sample_merge

8808c7c

afeldman-nm marked this pull request as ready for review February 11, 2025 15:21

afeldman-nm requested review from WoosukKwon, robertgshaw2-redhat and njhill as code owners February 11, 2025 15:21

afeldman-nm added 2 commits February 14, 2025 13:32

sampling params caching

39d3d0b

Signed-off-by: Andrew Feldman <afeldman@neuralmagic.com>

wip

5462e83

Signed-off-by: Andrew Feldman <afeldman@neuralmagic.com>

afeldman-nm added 21 commits February 16, 2025 21:21

wip

c0f8fb1

Signed-off-by: Andrew Feldman <afeldman@neuralmagic.com>

Merge branch 'main' into sample_merge

252e2ae

Merge branch 'main' into sample_merge

97feb43

parallel sampling core

103ceb6

Signed-off-by: Andrew Feldman <afeldman@neuralmagic.com>

wip

e6a1134

Signed-off-by: Andrew Feldman <afeldman@neuralmagic.com>

wip

625e161

Signed-off-by: Andrew Feldman <afeldman@neuralmagic.com>

add parallel sampling requests

e28388e

Signed-off-by: Andrew Feldman <afeldman@neuralmagic.com>

wip

1c18dc2

Signed-off-by: Andrew Feldman <afeldman@neuralmagic.com>

working parallel sampling in LLMEngine

c9c3dbb

Signed-off-by: Andrew Feldman <afeldman@neuralmagic.com>

Merge branch 'main' into sync_sample_merge

d927a4a

unit test

0f0075c

Signed-off-by: Andrew Feldman <afeldman@neuralmagic.com>

refactoring; bugfix

e4a0e6c

Signed-off-by: Andrew Feldman <afeldman@neuralmagic.com>

Merge branch 'main' into sync_sample_merge

bcdec42

seed

196fc68

Signed-off-by: Andrew Feldman <afeldman@neuralmagic.com>

refactor

f86708a

Signed-off-by: Andrew Feldman <afeldman@neuralmagic.com>

Merge branch 'main' into sample_merge

7e653f9

Parallel sampling unit tests

6b1be36

Signed-off-by: Andrew Feldman <afeldman@neuralmagic.com>

Merge branch 'main' into sync_sample_merge

081a695

Merge branch 'sync_sample_merge' into sync_sample

a3c5b22

Merge branch 'main' into sample_merge

13f3424

Merge branch 'sync_sample' into sample_merge

5f57964

afeldman-nm mentioned this pull request Feb 18, 2025

[V1] V1 engine implements parallel sampling, 2/2: LLMEngine support #13421

Closed

afeldman-nm changed the title ~~[V1] V1 engine implements parallel sampling, 1/2: AsyncLLM support~~ [V1] V1 engine implements parallel sampling (AsyncLLM and LLMEngine) Feb 18, 2025

pre-commit hook fixes

933a90e

Signed-off-by: Andrew Feldman <afeldman@neuralmagic.com>

njhill reviewed Feb 18, 2025

View reviewed changes

njhill reviewed Feb 19, 2025

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[V1] V1 engine implements parallel sampling (AsyncLLM and LLMEngine) #10980

[V1] V1 engine implements parallel sampling (AsyncLLM and LLMEngine) #10980

afeldman-nm commented Dec 7, 2024 •

edited by github-actions bot

Loading

github-actions bot commented Dec 7, 2024

mergify bot commented Dec 13, 2024

FrederickVu commented Jan 30, 2025

mergify bot commented Feb 5, 2025

m-harmonic commented Feb 14, 2025

afeldman-nm commented Feb 18, 2025

njhill left a comment

njhill Feb 18, 2025

njhill Feb 18, 2025

njhill Feb 18, 2025

njhill Feb 18, 2025

njhill Feb 18, 2025

njhill Feb 18, 2025

njhill Feb 18, 2025

njhill Feb 18, 2025

njhill Feb 18, 2025 •

edited

Loading

njhill Feb 19, 2025

njhill Feb 18, 2025

njhill Feb 19, 2025

njhill Feb 19, 2025

		async for out in merge_async_iterators(*gens):
		yield out[1] # out[0] is index

	if seed is None and self.cached_child_sampling_params:
	if self.cached_child_sampling_params:

	def _get_parent_request_output(self) -> RequestOutput:
	def _get_final_request_output(self) -> RequestOutput:

		from vllm.sampling_params import RequestOutputKind, SamplingParams


		class ParallelSamplingRequestManager:

	class ParallelSamplingRequestManager:
	class ParallelSamplingRequest:

		self._do_reset_parallel_sampling = (True if num_parallel_reqs > 0 else
		self._do_reset_parallel_sampling)

	self._do_reset_parallel_sampling = (True if num_parallel_reqs > 0 else
	self._do_reset_parallel_sampling)
	self._do_reset_parallel_sampling \|= num_parallel_reqs > 0

[V1] V1 engine implements parallel sampling (AsyncLLM and LLMEngine) #10980

Are you sure you want to change the base?

[V1] V1 engine implements parallel sampling (AsyncLLM and LLMEngine) #10980

Conversation

afeldman-nm commented Dec 7, 2024 • edited by github-actions bot Loading

github-actions bot commented Dec 7, 2024

mergify bot commented Dec 13, 2024

FrederickVu commented Jan 30, 2025

mergify bot commented Feb 5, 2025

m-harmonic commented Feb 14, 2025

afeldman-nm commented Feb 18, 2025

njhill left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

njhill Feb 18, 2025 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

afeldman-nm commented Dec 7, 2024 •

edited by github-actions bot

Loading

njhill Feb 18, 2025 •

edited

Loading